import string
from matplotlib.pylab import *

#read files generated from contextdecode
def read_MSE(filename):
   file = open(filename)
   while True:
       x = string.split(file.readline())
       if x[0] == "Outputting:":
           break
   return read_triplet_lines(file)

#generic read file f, assumed to numbers line by line. 
def read_triplet_lines(f):
   retval = []
   while True:
      x= f.readline()
      if x == "":
          break
      retval.append([float(y) for y in string.split(x)])
   return reorder(retval)

#(x1,y1,z2),...,(xn,yn,zn)-> (x1,..,xn),(y1,..,yn),(z1..zn)
def reorder(x):
    return [map(lambda x: x[i], x) for i in [0,1,2]]

def psnr(x):
    return 10.0*log10(255.0*255.0/float(x))

def plot_psnr(mse_triplet):
    bitrate = map(lambda x: x/512.0/512.0, mse_triplet[0])
    psnr1 = map(psnr, mse_triplet[1])
    psnr2 = map(psnr, mse_triplet[2])
    plot(bitrate, psnr1, '+')
    plot(bitrate, psnr2)

def plot_one_psnr(mse_triplet):
    bitrate = map(lambda x: x/512.0/512.0, mse_triplet[0])
    psnr2 = map(psnr, mse_triplet[2])
    plot(bitrate, psnr2)

def plot_one_psnr_var(mse_triplet, plottype):
    bitrate = map(lambda x: x/512.0/512.0, mse_triplet[0])
    psnr2 = map(psnr, mse_triplet[2])
    plot(bitrate, psnr2, plottype)

def plot_stat_bs_psnr(statfile, statbsfile):
    plot_one_psnr(read_MSE(statfile))
    plot_one_psnr_var(read_MSE(statbsfile), '+')

#read files generated from contextdecode
def plot_file(filename):
    plot_psnr(read_MSE(filename))
    axis([0, 1.4, 25, 42])
    show()

def plot_files(filename):
    plot_one_psnr(read_MSE(filename[0]))
    plot_one_psnr(read_MSE(filename[1]))
    axis([0, 1.4, 25, 42])
    show()

def plot_compare(fileprefix1, fileprefix2):
    plot_stat_bs_psnr(fileprefix1 + '.txt', fileprefix1+'bs.txt')
    plot_stat_bs_psnr(fileprefix2 + '.txt', fileprefix2+'bs.txt')
    axis([0, 1, 30, 40])
    show()


#plot_files(['goodstats/lennagtws4stat.txt', 'goodstats/lennagtwstat.txt'])
plot_stat_bs_psnr('goodstats/lennagtws4stats.txt', 'goodstats/lennagtws4statsbs.txt')
#plot_stat_bs_psnr('goodstats/lennagtwstat.txt', 'goodstats/lennagtwstatbs.txt')
#plot_stat_bs_psnr('../lenna1greedystat.txt', '../lenna1greedystatbs.txt')
#plot_stat_bs_psnr('../lennagtws6stats.txt', '../lennagtws6statsbs.txt')
#plot_stat_bs_psnr('../lennagtws7stats.txt', '../lennagtws7statsbs.txt')
#plot_stat_bs_psnr('../lenna1b15s7stats.txt', '../lenna1b15s7statsbs.txt')
#plot_stat_bs_psnr('../lenna1b15s8stats.txt', '../lenna1b15s8statsbs.txt')
axis([0, 1, 30, 40])
show()

#conclusions: sort 7 better than sort 6, (atleast on gtw)
